TRN2 Meshes and Configurations #916

apoorvtintin · 2025-01-10T00:48:35Z

This PR adds meshes for TRN2/1 for Fuji models and transformer layer configuration favorable to Neuron.

Neuron supports stacked transformer and GroupedQKVLinear instead of FusedGroupedQKVLinear for Grouped Query Attention (GQA)

This is a newer version of the PR #885. This PR resolved all comments and requested changes mentioned in the linked PR.

apoorvtintin · 2025-01-10T00:55:07Z

Added a ModelConfigModifier that overrides the class for a module. Allowing different model configurations based on Model size and platform.

kelvin-zou

Thank you for making such change, overall looks good. A few nit comments.

kelvin-zou · 2025-01-10T00:57:05Z

axlearn/common/trainer_config_modifier.py

+                continue
+            # Here we assume x.y.z format.
+            # One example would be model.decoder.transformer.layer.
+            target_modules = module_name.split(".")


Can you try to extract a common util function named something like
def replace_module_recursive(target_modules:str, config_key: str, target_config) and make it applied to both here and RematSpecModifier

I extracted a helper function, let me know if this looks good

axlearn/common/trainer_config_modifier_test.py

apoorvtintin · 2025-01-10T07:44:33Z

Added ParameterPartitionSpecModifier for parameters to shard Embeddings in a vocab parallel manner as described in Megatron LM.

axlearn/common/trainer_config_modifier.py

ruomingp · 2025-01-11T23:50:52Z

axlearn/common/trainer_config_modifier.py

+
+            found_module, parent_module, key_in_parent = find_target_module(module_name, cfg)
+
+            # Copy configurations from the config being replaced on a best effort basis


Wait, this behavior is not explained in the class comments. So we are not replacing but merging the configs? Maybe we should support a merge function instead?

Yeah the goal is to change the config to a similar module. This means most of the configuration can be reused from before. Essentially replacing the module but merging the config. Let me extract out a merge function.

Abstracted out a merge function let me know if more changes are needed for this.

apoorvtintin · 2025-01-12T07:08:47Z

@ruomingp Thank you for the review, I have addressed all your comments, please let me know if more changes are needed.

axlearn/common/trainer_config_modifier.py

ruomingp · 2025-01-13T06:35:53Z

axlearn/common/trainer_config_modifier.py

+        for module_name, model_cfg in self._model_cfg_modifications.items():
+            found_module = _find_target_module(module_name, cfg)


In utils.py we have get_recursively and set_recursively for Nested[...]. I wonder if it will be useful to add corresponding methods to ConfigBase. Then we can do something like:

Suggested change

for module_name, model_cfg in self._model_cfg_modifications.items():

found_module = _find_target_module(module_name, cfg)

for cfg_path, cfg_modification in self._model_cfg_modifications.items():

child_cfg = cfg.get_recursively(cfg_path)

child_cfg = cfg_modification(child_cfg, path=cfg_path)

cfg.set_recursively(cfg_path, value=child_cfg)

Added get_recursively and set_recursively functions to ConfigBase. Let me know if it looks good

I wonder if an alternative (which aims to simplify the ConfigBase api) is to do something similar to Python's sorted; we allow utils.get_resursively to take a value fn:

# Default behavior is to use key lookup: utils.get_recursively(..., value_fn=lambda k,v: v[k]) # Custom behavior can be attribute lookup: utils.get_recursively(..., value_fn=lambda k,v: getattr(v,k))

A benefit is that other non-config instances can also leverage get_recursively.

apoorvtintin · 2025-01-15T00:26:21Z

Added a more flexible PartitionSpecModifier that can modify multiple partition_spec attributes in a single module config.

kelvin-zou

Mostly lgtm, some minor comments.

axlearn/common/trainer_config_modifier.py

apoorvtintin · 2025-01-27T10:09:17Z

@ruomingp and @kelvin-zou thank you both for the review. I addressed all comments, please let me know if anymore changes are needed. PR looks clean now.

markblee · 2025-01-27T17:46:24Z

axlearn/common/config.py

+        key: str
+
+    def recursive_traverse(self, key_path: Sequence[str]) -> tuple[Any, str]:
+        """Recursively traverse the config to find the target key.


Please see other comment re get_recursively; also, I wonder whether we actually need recursion here (seems like a loop would be simpler).

markblee · 2025-01-27T17:47:21Z

axlearn/common/trainer_config_modifier.py

@@ -146,6 +137,110 @@ def __call__(self, cfg: SpmdTrainer.Config) -> SpmdTrainer.Config:
        return cfg


+class ModelConfigModifier(ConfigModifier):


Which part of this class is specific to model? It seems to take generic modifications?

markblee · 2025-01-27T17:47:46Z

axlearn/common/trainer_config_modifier.py

+        """Configure ModelConfigModifier.
+
+        Attributes:
+            model_cfg_modifications: A mapping from module path


markblee · 2025-01-27T17:48:15Z

axlearn/common/trainer_config_modifier.py

+        """Merge configurations from the config being replaced on a best effort basis.
+
+        Merge Rules:
+            - Klass is not changed, use target cfg


Suggested change

- Klass is not changed, use target cfg

- Klass is not changed, use target cfg.

Please end all sentences with punctuations.

markblee · 2025-01-27T17:48:36Z

axlearn/common/trainer_config_modifier.py

+            target_cfg: configuration that will replace found_module.
+            found_module: existing configuration whose class will be replaced


Suggested change

target_cfg: configuration that will replace found_module.

found_module: existing configuration whose class will be replaced

target_cfg: Configuration that will replace found_module.

found_module: Existing configuration whose class will be replaced

markblee · 2025-01-27T17:54:03Z

axlearn/experiments/text/gpt/fuji.py

+    if version != Version.V1:
+        trn2_model_modifications.append(
+            ModelConfigModifier.default_config().set(
+                target_config="model.decoder.transformer.layer.self_attention.attention."


A downside of representing these deeply nested configs as string paths is that they are brittle, and can quickly become outdated.

Have we considered using cfg.visit to achieve some of these modifications (e.g.,

axlearn/axlearn/common/layers.py

Line 266 in 1c22688

def set_layer_norm_eps_recursively(cfg: ConfigBase, eps: float, set_only_if_none: bool = False):

)?

(A bit late to review, so apologies if this discussion has already taken place.)

ruomingp · 2025-01-27T19:07:23Z

axlearn/common/config.py

+        # The key string.
+        key: str
+
+    def recursive_traverse(self, key_path: Sequence[str]) -> tuple[Any, str]:


Does this need to be a public method?

ruomingp · 2025-01-27T19:08:56Z

axlearn/common/config.py

+        traverse_result = self.recursive_traverse(key_path)
+        return getattr(traverse_result.parent, traverse_result.key)
+
+    def set_recursively(self, key_path: Sequence[str], new_value: Any):


Please name consistently with

axlearn/axlearn/common/utils.py

Lines 907 to 910 in a854738

def set_recursively(

x: NestedTensor,

*,

value: Tensor,

.

Suggested change

def set_recursively(self, key_path: Sequence[str], new_value: Any):

def set_recursively(self, path: Sequence[str], *, value: Any):

ruomingp · 2025-01-27T19:10:27Z

axlearn/common/config.py

+        value = getattr(self, target_key)
+        return value.recursive_traverse(key_path[1:])
+
+    def get_recursively(self, key_path: Sequence[str]) -> Any:


Suggested change

def get_recursively(self, key_path: Sequence[str]) -> Any:

def get_recursively(self, path: Sequence[str]) -> Any:

ruomingp · 2025-01-27T19:10:40Z

axlearn/common/config.py

+        """Recursively find the target key in the config and return its value.
+
+        Args:
+            key_path: A sequence of keys for indexing to get the target value.


Can path be empty? Maybe it can return self if path is empty?

ruomingp · 2025-01-27T19:10:53Z

axlearn/common/config.py

+        """Recursively find the target key in the config and set its value.
+
+        Args:
+            key_path: A sequence of keys for indexing to set the target value.


Can path be empty?

ruomingp · 2025-01-27T19:12:21Z

axlearn/common/config.py

+        Raises:
+            ValueError: A key in key_path is not found.
+        """
+        traverse_result = self.recursive_traverse(key_path)


Can we do something like:

Suggested change

traverse_result = self.recursive_traverse(key_path)

if not path:

raise ValueError(...)

parent = self.get_recursively(path[:-1])

apoorvtintin requested review from ruomingp, markblee and a team as code owners January 10, 2025 00:48

apoorvtintin mentioned this pull request Jan 10, 2025

Add meshes and config for TRN2/1 for Fuji models #885

Closed

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from 6b404f6 to 3f7c840 Compare January 10, 2025 00:53

kelvin-zou reviewed Jan 10, 2025

View reviewed changes

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from 708fc5e to d481132 Compare January 10, 2025 07:38

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from 5be50d7 to 9b10041 Compare January 10, 2025 08:10

apoorvtintin requested a review from kelvin-zou January 10, 2025 19:44

ruomingp requested changes Jan 11, 2025

View reviewed changes

apoorvtintin force-pushed the mainline-upstream-boilerplate branch from 9b10041 to 0f0a530 Compare January 12, 2025 07:06

apoorvtintin requested a review from ruomingp January 12, 2025 07:08

ruomingp requested changes Jan 13, 2025

View reviewed changes

apoorvtintin mentioned this pull request Jan 13, 2025

[DO-NOT-MERGE] PR encompassing all changes needed to support neuron on Axlearn #919

Open

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from c23e3b2 to 94bfff6 Compare January 15, 2025 00:10

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from 45c7df1 to 8807856 Compare January 17, 2025 01:17

kelvin-zou reviewed Jan 17, 2025

View reviewed changes

axlearn/common/trainer_config_modifier.py Outdated Show resolved Hide resolved

axlearn/common/trainer_config_modifier.py Outdated Show resolved Hide resolved

apoorvtintin force-pushed the mainline-upstream-boilerplate branch from 8807856 to 25510d6 Compare January 22, 2025 01:39

apoorvtintin requested a review from kelvin-zou January 22, 2025 01:40

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from eec33eb to 86bafa8 Compare January 23, 2025 05:40

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from 4661492 to fe96240 Compare January 23, 2025 16:38

TRN2 Meshes and Configurations

c9f9ab1

apoorvtintin force-pushed the mainline-upstream-boilerplate branch from fe96240 to da90757 Compare January 24, 2025 23:32

apoorvtintin requested a review from ruomingp January 24, 2025 23:40

apoorvtintin added 2 commits January 27, 2025 02:04

Add get_recursive and set_recursive to ConfigBase.

2225f87

Update golden configs

7e2e5f2

apoorvtintin force-pushed the mainline-upstream-boilerplate branch from da90757 to 7e2e5f2 Compare January 27, 2025 10:07

markblee reviewed Jan 27, 2025

View reviewed changes

ruomingp reviewed Jan 27, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

TRN2 Meshes and Configurations #916

TRN2 Meshes and Configurations #916

apoorvtintin commented Jan 10, 2025

apoorvtintin commented Jan 10, 2025

kelvin-zou left a comment

kelvin-zou Jan 10, 2025

apoorvtintin Jan 10, 2025 •

edited

Loading

apoorvtintin commented Jan 10, 2025 •

edited

Loading

ruomingp Jan 11, 2025

apoorvtintin Jan 12, 2025

apoorvtintin Jan 12, 2025

apoorvtintin commented Jan 12, 2025

ruomingp Jan 13, 2025

apoorvtintin Jan 27, 2025

markblee Jan 27, 2025

apoorvtintin commented Jan 15, 2025 •

edited

Loading

kelvin-zou left a comment

apoorvtintin commented Jan 27, 2025 •

edited

Loading

markblee Jan 27, 2025

markblee Jan 27, 2025

markblee Jan 27, 2025

markblee Jan 27, 2025

markblee Jan 27, 2025

markblee Jan 27, 2025

ruomingp Jan 27, 2025

ruomingp Jan 27, 2025

ruomingp Jan 27, 2025

ruomingp Jan 27, 2025

ruomingp Jan 27, 2025

ruomingp Jan 27, 2025


		found_module, parent_module, key_in_parent = find_target_module(module_name, cfg)

		# Copy configurations from the config being replaced on a best effort basis

		for module_name, model_cfg in self._model_cfg_modifications.items():
		found_module = _find_target_module(module_name, cfg)

-        for module_name, model_cfg in self._model_cfg_modifications.items():
-            found_module = _find_target_module(module_name, cfg)
+        for cfg_path, cfg_modification in self._model_cfg_modifications.items():
+            child_cfg = cfg.get_recursively(cfg_path)
+            child_cfg = cfg_modification(child_cfg, path=cfg_path)
+            cfg.set_recursively(cfg_path, value=child_cfg)

		@@ -146,6 +137,110 @@ def __call__(self, cfg: SpmdTrainer.Config) -> SpmdTrainer.Config:
		return cfg


		class ModelConfigModifier(ConfigModifier):

	- Klass is not changed, use target cfg
	- Klass is not changed, use target cfg.

		target_cfg: configuration that will replace found_module.
		found_module: existing configuration whose class will be replaced

	def set_recursively(self, key_path: Sequence[str], new_value: Any):
	def set_recursively(self, path: Sequence[str], *, value: Any):

	def get_recursively(self, key_path: Sequence[str]) -> Any:
	def get_recursively(self, path: Sequence[str]) -> Any:

-        traverse_result = self.recursive_traverse(key_path)
+        if not path:
+            raise ValueError(...)
+        parent = self.get_recursively(path[:-1])

TRN2 Meshes and Configurations #916

Are you sure you want to change the base?

TRN2 Meshes and Configurations #916

Conversation

apoorvtintin commented Jan 10, 2025

apoorvtintin commented Jan 10, 2025

kelvin-zou left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apoorvtintin Jan 10, 2025 • edited Loading

Choose a reason for hiding this comment

apoorvtintin commented Jan 10, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apoorvtintin commented Jan 12, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apoorvtintin commented Jan 15, 2025 • edited Loading

kelvin-zou left a comment

Choose a reason for hiding this comment

apoorvtintin commented Jan 27, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apoorvtintin Jan 10, 2025 •

edited

Loading

apoorvtintin commented Jan 10, 2025 •

edited

Loading

apoorvtintin commented Jan 15, 2025 •

edited

Loading

apoorvtintin commented Jan 27, 2025 •

edited

Loading